import asyncio
import websockets
import json
import base64
import time
from pydub import AudioSegment
from io import BytesIO

class ASRTestClient:
    """ASR测试客户端"""
    
    def __init__(self, uri: str = "ws://localhost:8888"):
        self.uri = uri
        
    async def test_realtime_mode(self):
        """测试实时识别模式"""
        print("测试实时识别模式...")
        
        async with websockets.connect(self.uri) as websocket:
            # 选择实时模式
            await websocket.send(json.dumps({
                "type": "select_mode",
                "mode": "realtime"
            }))
            
            response = await websocket.recv()
            print(f"服务器响应: {response}")
            
            # 模拟发送音频数据
            for i in range(5):
                start_time = 4.5*i
                end_time = start_time+4.5
                # 创建模拟音频数据
                audio_b64 = b''
                dummy_audio = b'\x00' * 1024  # 1KB的静音数据
                #audio_b64 = base64.b64encode(dummy_audio).decode()
                
                # with open("testRes/input.wav", 'rb') as audio_file:
                #     audio_data = audio_file.read()
                #     audio_b64 = base64.b64encode(audio_data).decode()

                audio = AudioSegment.from_file("testRes/input.wav")
                # 转换为毫秒
                start_ms = start_time * 1000
                end_ms = end_time * 1000

                # 截取音频片段
                segment = audio[start_ms:end_ms]

                # 创建内存缓冲区
                buffer = BytesIO()

                # 将音频片段导出到缓冲区
                segment.export(buffer, format="wav")

                # 获取字节数组
                audio_data = buffer.getvalue()
                audio_b64 = base64.b64encode(audio_data).decode()

                await websocket.send(json.dumps({
                    "type": "audio_chunk",
                    "audio": audio_b64
                }))
                
                print(f"发送音频块 {i+1}")
                await asyncio.sleep(1)  # 模拟实时音频流
                
                # 接收可能的识别结果
                try:
                    response = await asyncio.wait_for(websocket.recv(), timeout=0.1)
                    print(f"识别结果: {response}")
                except asyncio.TimeoutError:
                    pass
            
            # 结束音频流
            await websocket.send(json.dumps({
                "type": "end_audio"
            }))
            
            # 接收最终结果
            response = await websocket.recv()
            print(f"最终响应: {response}")
    
    async def test_file_mode(self):
        """测试文件识别模式"""
        print("测试文件识别模式...")
        
        async with websockets.connect(self.uri) as websocket:
            # 选择文件模式
            await websocket.send(json.dumps({
                "type": "select_mode",
                "mode": "file"
            }))
            
            response = await websocket.recv()
            print(f"服务器响应: {response}")
            
            # 创建模拟音频文件数据
            file_b64 = b''
            dummy_file = b'\x00' * 8192  # 8KB的模拟音频文件
            #file_b64 = base64.b64encode(dummy_file).decode()

            with open("testRes/input.wav", 'rb') as audio_file:
                audio_data = audio_file.read()
                file_b64 = base64.b64encode(audio_data).decode()

            await websocket.send(json.dumps({
                "type": "file_upload",
                #"filename": "test_audio.wav",
                "file_data": file_b64
            }))
            
            print("已上传音频文件")
            
            # 接收识别结果
            response = await websocket.recv()
            print(f"文件识别结果: {response}")

async def main():
    """主测试函数"""
    client = ASRTestClient()
    
    print("=== ASR服务器测试客户端 ===")
    print("请确保ASR服务器已启动 (python asr_server.py)")
    
    try:
        # 测试实时模式
        await client.test_realtime_mode()
        print("\n" + "="*50 + "\n")
        
        # 测试文件模式
        await client.test_file_mode()
        
    except ConnectionRefusedError:
        print("错误: 无法连接到ASR服务器，请确保服务器已启动")
    except Exception as e:
        print(f"测试过程中发生错误: {e}")

if __name__ == "__main__":
    asyncio.run(main())